import paho.mqtt.client as mqtt
import requests
import json
import time
import logging
import argparse  # Add this import at the top
import threading
import os
import psutil
import netifaces
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

# Move logging configuration into a function
def configure_logging(debug_mode=False):
    """Configure logging based on debug mode"""
    if debug_mode:
        # Log to terminal in debug mode
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
    else:
        # Log to daily file in production mode
        log_dir = os.path.dirname(os.path.abspath(__file__))
        current_date = datetime.now().strftime('%Y-%m-%d')
        log_file = os.path.join(log_dir, f'mqtt-{current_date}.log')
        
        # Create formatter
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        
        # Create file handler
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(formatter)
        
        # Create console handler
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        
        # Get logger and set level
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        
        # Remove any existing handlers
        logger.handlers = []
        
        # Add handlers
        logger.addHandler(file_handler)
        logger.addHandler(console_handler)

# Initialize logger
logger = logging.getLogger(__name__)

class MQTTClient:
    # Class constants
    STATUS_UPDATE_INTERVAL = 30  # seconds
    RECONNECT_DELAY = 5  # seconds
    CONNECTION_TIMEOUT = 60  # seconds

    def __init__(self, cli_sn=None):
        # Load config and configure logging
        try:
            with open('config.json', 'r') as f:
                config = json.load(f)
            debug_mode = config.get('debug', False)
            configure_logging(debug_mode)
        except Exception as e:
            # Default to console logging if config file cannot be read
            configure_logging(True)
            logger.error(f"Error reading config file: {e}, defaulting to debug mode logging")

        self.executor = ThreadPoolExecutor(max_workers=5)  # Add thread pool
        self.pending_tests = {} 
        self.device_sn = None
        self.cli_sn = cli_sn
        self.mqtt_client = None
        self.is_logged_in = False
        self.version = "1.0.0"
        self.running = True
        self.login_retry_timer = None  # Add this line
        self.login_retry_interval = 5  # Default value
        self.message_seq = 0  # Add message sequence counter

    def _execute_speed_test(self, data, request_seq):
        """Execute speed test in background thread"""
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            result = None
            
            # If in debug mode, simulate test delay and read from test file
            if config.get('debug', False):
                logger.info("Debug mode: Simulating speed test for 30 seconds")
                time.sleep(30)
                
                try:
                    with open('test_speed_test.json', 'r') as f:
                        result = {
                            "code": 0,
                            "msg": "success",
                            "result": json.load(f)
                        }
                except FileNotFoundError:
                    result = {
                        "code": 1,
                        "msg": "Debug test file not found",
                        "result": {"test_status": "error"}
                    }
            else:
                # Perform actual speed test
                headers = {'Content-Type': 'application/json'}
                url = 'http://127.0.0.1:12789/v1/recruit_detector/speed_test'
                
                response = requests.post(url, headers=headers, json=data)
                
                if response.status_code == 200:
                    result = {
                        "code": 0,
                        "msg": "success",
                        "result": response.json()
                    }
                else:
                    result = {
                        "code": 1,
                        "msg": f"Speed test failed: HTTP {response.status_code}",
                        "result": {"test_status": "failed"}
                    }
            
            # Send response when test completes
            response = {
                "cmd": "route_test_reply",
                "sn": self.device_sn,
                "seq": request_seq,
                "timestamp": int(time.time()),
                "data": result
            }
            self.mqtt_client.publish(f"{self.device_sn}s", json.dumps(response))
            
        except Exception as e:
            logger.error(f"Error in speed test execution: {e}")
            error_response = {
                "cmd": "route_test_reply",
                "sn": self.device_sn,
                "seq": request_seq,
                "timestamp": int(time.time()),
                "data": {
                    "code": 1,
                    "msg": f"Speed test error: {str(e)}",
                    "result": {"test_status": "error"}
                }
            }
            self.mqtt_client.publish(f"{self.device_sn}s", json.dumps(error_response))
        finally:
            # Clean up tracking
            if request_seq in self.pending_tests:
                del self.pending_tests[request_seq]
    
    def get_next_seq(self):
        """Get next message sequence number"""
        self.message_seq += 1
        return self.message_seq

    def get_mqtt_config(self):
        """Get MQTT server configuration from remote API or debug config"""
        try:
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            required_fields = ['host', 'port']
            
            if config.get('debug', False):
                debug_config = config.get('debug_mqtt_config', {})
                if not all(field in debug_config for field in required_fields):
                    raise ValueError("Missing required fields in debug MQTT config")
                return debug_config
            
            api_url = config.get('api', {}).get('mqtt_config_url')
            retry_interval = config.get('api', {}).get('config_retry_interval', 5)

            if not api_url:
                raise ValueError("Missing MQTT config URL in configuration")
                
            while True:
                try:
                    response = requests.get(api_url)
                    if response.status_code == 200:
                        response_data = response.json()
                        
                        # Check if response has the expected structure
                        if response_data.get('code') == 200 and isinstance(response_data.get('data'), dict):
                            mqtt_config = response_data['data']
                            
                            # Convert port to integer if it's a string
                            if isinstance(mqtt_config.get('port'), str):
                                mqtt_config['port'] = int(mqtt_config['port'])
                                
                            if not all(field in mqtt_config for field in required_fields):
                                raise ValueError("Missing required fields in MQTT config from API")
                                
                            logger.info("Successfully retrieved MQTT config from API")
                            return mqtt_config
                        else:
                            logger.error("Invalid response format from API")
                    else:
                        logger.error(f"Failed to get MQTT config: HTTP {response.status_code}")
                except requests.RequestException as e:
                    logger.error(f"Error accessing MQTT config API: {e}")
                except json.JSONDecodeError as e:
                    logger.error(f"Invalid JSON response from API: {e}")
                except Exception as e:
                    logger.error(f"Unexpected error getting MQTT config: {e}")
                
                logger.info(f"Retrying MQTT config retrieval in {retry_interval} seconds...")
                time.sleep(retry_interval)
                
        except Exception as e:
            logger.error(f"Error reading configuration file: {e}")
            return None
        
    def get_device_sn(self):
        """Get device SN code from CLI arg, config or file"""
        # First check if SN was provided via CLI
        if self.cli_sn:
            logger.info("Using SN from command line argument")
            return self.cli_sn

        try:
            # Load configuration from config file
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Check if debug mode is enabled
            if config.get('debug', False):
                logger.info("Debug mode: Using debug device SN")
                return config['debug_device_sn']
            
            # If not in debug mode, get SN from file
            sn_file = '/etc/app/.sn'
            with open(sn_file, 'r') as f:
                return f.read().strip()
        except Exception as e:
            logger.error(f"Error getting device SN: {e}")
            return None

    def on_connect(self, client, userdata, flags, rc):
        """Callback when connected to MQTT broker"""
        logger.info(f"Connected with result code {rc}")
        if rc == 0:
            # Subscribe to the SN-specific topic
            client.subscribe(f"{self.device_sn}")
            # Start login process
            self.start_login_process()

    def process_get_eth(self, data):
        """Process get ethernet interfaces command"""
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            if config.get('debug', False):
                # Return mock data in debug mode
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {
                        "interfaces": [
                            {"name": "eth0", "status": "up", "mac": "00:11:22:33:44:55"},
                            {"name": "eth1", "status": "down", "mac": "66:77:88:99:AA:BB"}
                        ]
                    }
                }
            else:
                # Get real system network interfaces
                
                interfaces = []
                for iface in psutil.net_if_addrs().keys():
                    if iface.startswith('eth') or iface.startswith('en'):
                        try:
                            mac = netifaces.ifaddresses(iface)[netifaces.AF_LINK][0]['addr']
                            status = "up" if iface in psutil.net_if_stats() and psutil.net_if_stats()[iface].isup else "down"
                            interfaces.append({
                                "name": iface,
                                "status": status,
                                "mac": mac
                            })
                        except (KeyError, IndexError) as e:
                            logger.error(f"Error getting info for interface {iface}: {e}")
                            continue
                
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {"interfaces": interfaces}
                }
                
        except Exception as e:
            logger.error(f"Error getting ethernet interfaces: {e}")
            return {
                "code": 1,
                "msg": f"Failed to get ethernet interfaces: {str(e)}",
                "result": {"interfaces": []}
            }
    
    def start_login_process(self):
        """Start the login process with retry mechanism"""
        try:
            # Read config to get retry interval
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            self.login_retry_interval = config.get('login_retry_interval', 5)
            
            # Send initial login request
            self.send_login_request()
            
            # Start retry timer
            self.schedule_login_retry()
            
        except Exception as e:
            logger.error(f"Error starting login process: {e}")

    def schedule_login_retry(self):
        """Schedule the next login retry if not logged in"""
        if not self.is_logged_in and self.running:
            self.login_retry_timer = threading.Timer(self.login_retry_interval, self.send_login_request)
            self.login_retry_timer.start()

    def send_login_request(self):
        """Send login request after connection"""
        if self.is_logged_in:
            return

        try:
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            login_topic = config.get('login_topic', f"{self.device_sn}s")

            login_msg = {
                "cmd": "login",
                "sn": self.device_sn,
                "seq": self.get_next_seq(),  # Add sequence number
                "timestamp": int(time.time()),
                "data": {
                    "version": self.version
                }
            }
            self.mqtt_client.publish(login_topic, json.dumps(login_msg))
            logger.info(f"Login request sent to topic: {login_topic}")
            
            self.schedule_login_retry()
            
        except Exception as e:
            logger.error(f"Error sending login request: {e}")
            try:
                self.mqtt_client.publish(f"{self.device_sn}s", json.dumps(login_msg))
                logger.info("Login request sent to fallback topic")
                self.schedule_login_retry()
            except Exception as e:
                logger.error(f"Error sending login request to fallback topic: {e}")

    def process_config_route(self, data):
        """Process route configuration command"""
        try:
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Determine output path based on debug mode
            output_path = 'test_config.json' if config.get('debug', False) else '/app/system/edge_router/conf/config.json'
            
            # Write the configuration data
            try:
                with open(output_path, 'w') as f:
                    json.dump(data, f, indent=2)
                logger.info(f"Route configuration written to {output_path}")
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {"status": "configured"}
                }
            except Exception as e:
                logger.error(f"Failed to write configuration: {e}")
                return {
                    "code": 1,
                    "msg": f"Failed to write configuration: {str(e)}",
                    "result": {"status": "failed"}
                }
                
        except Exception as e:
            logger.error(f"Error reading debug configuration: {e}")
            return {
                "code": 1,
                "msg": f"Configuration error: {str(e)}",
                "result": {"status": "failed"}
            }

    def process_get_route_status(self, data):
        """Process get route status command"""
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Determine which status file to read based on debug mode
            status_file = 'test_multidialstatus.json' if config.get('debug', False) else '/tmp/line_status.json'
            
            try:
                with open(status_file, 'r') as f:
                    status_data = json.load(f)
                    return {
                        "code": 0,
                        "msg": "success",
                        "result": status_data
                    }
            except FileNotFoundError:
                logger.error(f"Status file not found: {status_file}")
                return {
                    "code": 1,
                    "msg": f"Status file not found: {status_file}",
                    "result": {"status": "unknown"}
                }
            except json.JSONDecodeError as e:
                logger.error(f"Invalid JSON in status file: {e}")
                return {
                    "code": 1,
                    "msg": "Invalid status file format",
                    "result": {"status": "error"}
                }
                
        except Exception as e:
            logger.error(f"Error reading route status: {e}")
            return {
                "code": 1,
                "msg": f"Failed to get route status: {str(e)}",
                "result": {"status": "error"}
            }
        
    def process_boot_cloud_platform(self, data):
        """
        Process boot_cloud_platform command to launch cloud platform docker containers
        
        Expected data format:
        {
            "id": 1,
            "name": "七牛云",
            "image_name": "vm_3rd_niulink:latest",
            "container_name": "vm-niulink"
        }
        """
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Extract container and image information
            container_name = data.get("container_name")
            image_name = data.get("image_name")
            
            if not container_name or not image_name:
                logger.error("Missing required container_name or image_name in boot request")
                return {
                    "code": 1,
                    "msg": "Missing required parameters",
                    "result": {"status": "failed"}
                }
            
            # In debug mode, simulate the operation
            if config.get('debug', False):
                logger.info(f"Debug mode: Simulating boot of {image_name} as {container_name}")
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {"status": "started"}
                }
            
            # First check if the container is already running
            import subprocess
            check_cmd = ["docker", "-H", "unix:///var/run/docker-proxy.sock", "ps", "--filter", f"name={container_name}"]
            
            logger.info(f"Checking if container {container_name} is already running")
            check_process = subprocess.run(
                check_cmd,
                capture_output=True,
                text=True
            )
            
            # If container name appears in the output, it's already running
            if container_name in check_process.stdout:
                logger.info(f"Container {container_name} is already running")
                return {
                    "code": 0,
                    "msg": "Container is already running",
                    "result": {"status": "started"}
                }
            
            # Container is not running, so start it
            logger.info(f"Container {container_name} not found, starting it now")
            docker_cmd = [
                "docker", "-H", "unix:///var/run/docker-proxy.sock", 
                "run", "-itd", "--name", container_name, image_name
            ]
            
            # Execute the docker run command
            logger.info(f"Executing: {' '.join(docker_cmd)}")
            process = subprocess.run(
                docker_cmd,
                capture_output=True,
                text=True
            )
            
            if process.returncode != 0:
                logger.error(f"Failed to start container: {process.stderr}")
                return {
                    "code": 1,
                    "msg": f"Failed to start container: {process.stderr}",
                    "result": {"status": "failed"}
                }
            
            # Check if the container is actually running after starting
            verify_cmd = ["docker", "-H", "unix:///var/run/docker-proxy.sock", "ps", "--filter", f"name={container_name}"]
            verify_process = subprocess.run(
                verify_cmd,
                capture_output=True,
                text=True
            )
            
            if container_name in verify_process.stdout:
                logger.info(f"Container {container_name} successfully started")
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {"status": "started"}
                }
            else:
                logger.error(f"Container {container_name} not found in running containers after starting")
                return {
                    "code": 1,
                    "msg": "Container not found after starting",
                    "result": {"status": "failed"}
                }
                
        except Exception as e:
            logger.error(f"Error booting cloud platform: {e}")
            return {
                "code": 1,
                "msg": f"Failed to boot cloud platform: {str(e)}",
                "result": {"status": "failed"}
            }
        
    def process_get_nic_status(self, data):
        """Process get network interface card status command"""
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Determine which status file to read based on debug mode
            status_file = 'test_nic_status.json' if config.get('debug', False) else '/tmp/nic_status.json'
            
            try:
                with open(status_file, 'r') as f:
                    status_data = json.load(f)
                    return {
                        "code": 0,
                        "msg": "success",
                        "result": status_data
                    }
            except FileNotFoundError:
                logger.error(f"NIC status file not found: {status_file}")
                return {
                    "code": 1,
                    "msg": f"NIC status file not found: {status_file}",
                    "result": {"status": "unknown"}
                }
            except json.JSONDecodeError as e:
                logger.error(f"Invalid JSON in NIC status file: {e}")
                return {
                    "code": 1,
                    "msg": "Invalid NIC status file format",
                    "result": {"status": "error"}
                }
                
        except Exception as e:
            logger.error(f"Error reading NIC status: {e}")
            return {
                "code": 1,
                "msg": f"Failed to get NIC status: {str(e)}",
                "result": {"status": "error"}
            }
    
    def process_get_platform_qrcode(self, data):
        """
        Process get_platform_qrcode command to retrieve QR code from cloud platform container
        
        Expected data format:
        {
            "id": 1,
            "name": "七牛云",
            "image_name": "vm_3rd_niulink:latest",
            "container_name": "vm-niulink"
        }
        """
        try:
            # Read configuration to check debug mode
            with open('config.json', 'r') as f:
                config = json.load(f)
            
            # Extract container name
            container_name = data.get("container_name")
            
            if not container_name:
                logger.error("Missing required container_name in QR code request")
                return {
                    "code": 1,
                    "msg": "Missing container_name parameter",
                    "result": {"status": "failed"}
                }
            
            # In debug mode, simulate the operation
            if config.get('debug', False):
                logger.info(f"Debug mode: Simulating QR code retrieval from {container_name}")
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {
                        "status": "success",
                        "qrcode": "https://example.com/debugqrcode"
                    }
                }
            
            # Construct the docker exec command
            import subprocess
            docker_cmd = [
                "docker", "-H", "unix:///var/run/docker-proxy.sock", 
                "exec", "-i", container_name, "virt_screen_qrcode", "--sn"
            ]
            
            # Execute the docker exec command
            logger.info(f"Executing: {' '.join(docker_cmd)}")
            process = subprocess.run(
                docker_cmd,
                capture_output=True,
                text=True
            )
            
            # Get the QR code from command output regardless of return code
            qrcode_data = process.stdout.strip()
            
            if qrcode_data:
                logger.info(f"Successfully retrieved QR code from container {container_name}")
                return {
                    "code": 0,
                    "msg": "success",
                    "result": {
                        "status": "success",
                        "qrcode": qrcode_data
                    }
                }
            else:
                # If stdout is empty, check if there's any error message
                error_msg = process.stderr.strip() if process.stderr else "Unknown error"
                logger.error(f"QR code command returned empty result: {error_msg}")
                return {
                    "code": 1,
                    "msg": f"Failed to get QR code: {error_msg}",
                    "result": {"status": "failed"}
                }
                    
        except Exception as e:
            logger.error(f"Error getting platform QR code: {e}")
            return {
                "code": 1,
                "msg": f"Failed to get QR code: {str(e)}",
                "result": {"status": "failed"}
            }

    def process_route_test(self, data, request_seq):  # Add request_seq parameter
        """Process route test command asynchronously"""
        try:
            # Start test in background using the original request sequence number
            self.pending_tests[request_seq] = self.executor.submit(
                self._execute_speed_test, data, request_seq
            )
            
            # Return immediate acknowledgment with original sequence number
            return {
                "code": 0,
                "msg": "Test started",
                "result": {"test_status": "running"}
            }
            
        except Exception as e:
            logger.error(f"Failed to start speed test: {e}")
            return {
                "code": 1,
                "msg": f"Failed to start test: {str(e)}",
                "result": {"test_status": "error"}
            }
    
    def on_disconnect(self, client, userdata, rc):
        """Callback when disconnected from MQTT broker"""
        logger.warning(f"Disconnected with result code: {rc}")
        self.is_logged_in = False
        
        # Reconnect if disconnection was unexpected
        if rc != 0 and self.running:
            logger.info("Attempting to reconnect...")
            while self.running:
                try:
                    self.mqtt_client.reconnect()
                    break
                except Exception as e:
                    logger.error(f"Reconnection failed: {e}")
                    time.sleep(self.RECONNECT_DELAY)

    def on_message(self, client, userdata, msg):
        """Callback when message is received"""
        try:
            payload = json.loads(msg.payload.decode())
            logger.info(f"Received message on {msg.topic}: {payload}")

            cmd = payload.get("cmd")
            response_data = None
            request_seq = payload.get("seq")  # Get sequence number from request

            if cmd == "login_reply" and not self.is_logged_in:
                if payload["data"]["code"] == 0:
                    self.is_logged_in = True
                    logger.info("Login successful")
                else:
                    logger.error(f"Login failed: {payload['data']['msg']}")

            elif self.is_logged_in:
                if cmd == "config_route":
                    response_data = self.process_config_route(payload["data"])
                elif cmd == "get_route_status":
                    response_data = self.process_get_route_status(payload["data"])
                elif cmd == "route_test":
                    response_data = self.process_route_test(payload["data"], request_seq)  # Pass request_seq
                elif cmd == "get_eth":  # Add new command handler
                    response_data = self.process_get_eth(payload.get("data", {}))
                elif cmd == "get_nic_status":  # Add new command handler
                    response_data = self.process_get_nic_status(payload.get("data", {}))
                elif cmd == "boot_cloud_platform":  # Add new cloud platform boot handler
                    response_data = self.process_boot_cloud_platform(payload.get("data", {}))
                elif cmd == "get_platform_qrcode":  # Add new QR code handler
                    response_data = self.process_get_platform_qrcode(payload.get("data", {}))

                if response_data:
                    response = {
                        "cmd": f"{cmd}_reply",
                        "sn": self.device_sn,
                        "seq": request_seq,  # Use original request sequence
                        "timestamp": int(time.time()),
                        "data": response_data
                    }
                    self.mqtt_client.publish(f"{self.device_sn}s", json.dumps(response))

        except Exception as e:
            logger.error(f"Error processing message: {e}")

    def run(self):
        """Main method to run the MQTT client"""
        # Get device SN
        self.device_sn = self.get_device_sn()
        if not self.device_sn:
            logger.error("Failed to get device SN")
            return

        # Get MQTT configuration
        mqtt_config = self.get_mqtt_config()
        if not mqtt_config:
            logger.error("Failed to get MQTT configuration")
            return

        # Initialize MQTT client
        self.mqtt_client = mqtt.Client(client_id=self.device_sn)  # Use device_sn as client_id
        self.mqtt_client.on_connect = self.on_connect
        self.mqtt_client.on_message = self.on_message
        self.mqtt_client.on_disconnect = self.on_disconnect

        # Set credentials if required
        if 'username' in mqtt_config and 'password' in mqtt_config:
            self.mqtt_client.username_pw_set(mqtt_config['username'], mqtt_config['password'])

        try:
            # Connect to MQTT broker
            self.mqtt_client.connect(mqtt_config['host'], mqtt_config['port'], self.CONNECTION_TIMEOUT)
            self.mqtt_client.loop_start()

            # Main loop
            while self.running:
                if self.is_logged_in:
                    with open('config.json', 'r') as f:
                        config = json.load(f)
                    heartbeat_topic = config.get('heartbeat_topic', 'device/heartbeat')

                    status_msg = {
                        "cmd": "status_update",
                        "sn": self.device_sn,
                        "seq": self.get_next_seq(),  # Add sequence number
                        "timestamp": int(time.time()),
                        "data": {
                            "status": "running"
                        }
                    }
                    self.mqtt_client.publish(heartbeat_topic, json.dumps(status_msg))
                time.sleep(self.STATUS_UPDATE_INTERVAL)

        except KeyboardInterrupt:
            logger.info("Stopping...")
            self.running = False
        except Exception as e:
            logger.error(f"Error in main loop: {e}")
        finally:
            self.cleanup()

    def cleanup(self):
        """Clean up resources and ensure proper shutdown"""
        self.running = False
        
        # Shutdown thread pool
        if hasattr(self, 'executor'):
            self.executor.shutdown(wait=False)

        # Cancel any pending login retry
        if self.login_retry_timer:
            self.login_retry_timer.cancel()
        
        if self.mqtt_client:
            try:
                # Send offline status before disconnecting
                status_msg = {
                    "cmd": "status_update",
                    "sn": self.device_sn,
                    "seq": self.get_next_seq(),  # Add sequence number
                    "timestamp": int(time.time()),
                    "data": {
                        "status": "offline"
                    }
                }
                with open('config.json', 'r') as f:
                        config = json.load(f)
                heartbeat_topic = config.get('heartbeat_topic', 'device/heartbeat')

                self.mqtt_client.publish(heartbeat_topic, json.dumps(status_msg))
                self.mqtt_client.loop_stop()
                self.mqtt_client.disconnect()
            except Exception as e:
                logger.error(f"Error during cleanup: {e}")

if __name__ == "__main__":
    # Set up argument parser
    parser = argparse.ArgumentParser(description='MQTT Client')
    parser.add_argument('--sn', type=str, help='Device serial number')
    args = parser.parse_args()

    # Initialize client with command line SN if provided
    mqtt_client = MQTTClient(cli_sn=args.sn)
    mqtt_client.run()